{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Semantic Kitti vizualisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "import sys\n",
    "import panel as pn\n",
    "import numpy as np\n",
    "import pyvista as pv\n",
    "import glob\n",
    "from matplotlib.colors import ListedColormap\n",
    "from omegaconf import OmegaConf\n",
    "\n",
    "pn.extension('vtk')\n",
    "os.system('/usr/bin/Xvfb :99 -screen 0 1024x768x24 &')\n",
    "os.environ['DISPLAY'] = ':99'\n",
    "os.environ['PYVISTA_OFF_SCREEN'] = 'True'\n",
    "os.environ['PYVISTA_USE_PANEL'] = 'True'\n",
    "pv.set_plot_theme(\"document\")\n",
    "\n",
    "from torch_points3d.datasets.segmentation.semantickitti import SemanticKittiDataset, SemanticKitti\n",
    "from torch_points3d.datasets.segmentation import IGNORE_LABEL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Scannet dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DIR = os.path.dirname(os.getcwd())\n",
    "dataset_options = OmegaConf.load(os.path.join(DIR,'conf/data/segmentation/semanticKitti.yaml'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_options.data.dataroot = os.path.join(DIR,\"data\")\n",
    "dataset = SemanticKittiDataset(dataset_options.data)\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualise the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def buil_cmap():\n",
    "    mapping = np.linspace(IGNORE_LABEL, 20, 256)\n",
    "    newcolors = np.zeros((256, 3))\n",
    "    for i, raw_label in SemanticKitti.LEARNING_MAP_INV.items():\n",
    "        color = SemanticKitti.COLOR_MAP[raw_label]\n",
    "        newcolors[mapping >= i] = np.asarray(color) / 255.\n",
    "    return ListedColormap(newcolors)\n",
    "cmap = buil_cmap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_random_data(event):\n",
    "    i = np.random.randint(0, len(dataset.train_dataset))\n",
    "    sample = dataset.train_dataset[i]\n",
    "    xyz = sample.pos.numpy()\n",
    "    pl = pv.Plotter(notebook=True)\n",
    "    point_cloud = pv.PolyData(xyz)\n",
    "    point_cloud['label'] = sample.y.numpy()\n",
    "    pl.add_points(point_cloud,cmap=cmap) \n",
    "    \n",
    "    pan.object = pl.ren_win\n",
    "    return point_cloud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl = pv.Plotter(notebook=True)\n",
    "pan = pn.panel(pl.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=True,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "button = pn.widgets.Button(name='New Random', button_type='primary')\n",
    "button.on_click(load_random_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dashboard = pn.Row(\n",
    "    pn.Column('## Visualise sample',button),\n",
    "    pan\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dashboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
